from functools import wraps
from flask import session, request, jsonify, g
from extensions import db
from models import User, RequestRecord, CostRecord, ModelPricing
from datetime import datetime
import jwt
import os

# JWT配置
JWT_SECRET_KEY = os.environ.get('SECRET_KEY', 'biomni-web-interface-secret-key-change-in-production')
JWT_ALGORITHM = 'HS256'
JWT_EXPIRATION_HOURS = 24

def generate_token(user_id):
    """生成JWT token"""
    from datetime import datetime, timedelta, timezone
    now = datetime.now(timezone.utc)
    payload = {
        'user_id': user_id,
        'exp': now + timedelta(hours=JWT_EXPIRATION_HOURS),
        'iat': now
    }
    return jwt.encode(payload, JWT_SECRET_KEY, algorithm=JWT_ALGORITHM)

def verify_token(token):
    """验证JWT token"""
    try:
        payload = jwt.decode(token, JWT_SECRET_KEY, algorithms=[JWT_ALGORITHM])
        return payload['user_id']
    except jwt.ExpiredSignatureError:
        return None  # token过期
    except jwt.InvalidTokenError:
        return None  # token无效

def login_required(f):
    """登录验证装饰器"""
    @wraps(f)
    def decorated_function(*args, **kwargs):
        # 检查session中的用户ID
        user_id = session.get('user_id')
        
        # 如果session中没有，检查Authorization header
        if not user_id:
            auth_header = request.headers.get('Authorization')
            if auth_header and auth_header.startswith('Bearer '):
                token = auth_header.split(' ')[1]
                user_id = verify_token(token)
        
        if not user_id:
            return jsonify({
                'success': False,
                'message': '请先登录',
                'error_code': 'UNAUTHORIZED'
            }), 401
        
        # 获取用户信息
        user = User.query.get(user_id)
        if not user or not user.is_active:
            return jsonify({
                'success': False,
                'message': '用户不存在或已被禁用',
                'error_code': 'USER_INACTIVE'
            }), 401
        
        # 将用户信息存储到g对象中，供视图函数使用
        g.current_user = user
        return f(*args, **kwargs)
    
    return decorated_function

def admin_required(f):
    """管理员权限验证装饰器"""
    @wraps(f)
    @login_required
    def decorated_function(*args, **kwargs):
        if g.current_user.role.value != 'admin':
            return jsonify({
                'success': False,
                'message': '需要管理员权限',
                'error_code': 'ADMIN_REQUIRED'
            }), 403
        
        return f(*args, **kwargs)
    
    return decorated_function

def request_limit_check(f):
    """请求限制检查装饰器"""
    @wraps(f)
    @login_required
    def decorated_function(*args, **kwargs):
        user = g.current_user
        
        # 检查用户是否可以发起请求
        can_request, message = user.can_make_request()
        if not can_request:
            return jsonify({
                'success': False,
                'message': message,
                'error_code': 'REQUEST_LIMIT_EXCEEDED',
                'weekly_count': user.get_weekly_request_count(),
                'weekly_limit': user.weekly_request_limit
            }), 429  # Too Many Requests
        
        return f(*args, **kwargs)
    
    return decorated_function

def create_request_record(user_id, session_id, query, model):
    """创建请求记录"""
    request_record = RequestRecord(
        user_id=user_id,
        session_id=session_id,
        query=query,
        model=model,
        status='pending',
        started_at=datetime.utcnow()
    )
    
    db.session.add(request_record)
    db.session.commit()
    
    return request_record

def update_request_record(request_id, status, response=None, error_message=None, 
                         input_tokens=0, output_tokens=0):
    """更新请求记录"""
    request_record = db.session.query(RequestRecord).get(request_id)
    if not request_record:
        return None
    
    request_record.status = status
    request_record.completed_at = datetime.utcnow()
    request_record.calculate_duration()
    
    if response:
        request_record.response = response
    if error_message:
        request_record.error_message = error_message
    
    # 更新token统计
    request_record.input_tokens = input_tokens
    request_record.output_tokens = output_tokens
    request_record.total_tokens = input_tokens + output_tokens
    
    # 如果请求成功完成，创建费用记录
    if status == 'completed' and (input_tokens > 0 or output_tokens > 0):
        create_cost_record(request_record)
    
    db.session.commit()
    return request_record

def create_cost_record(request_record):
    """创建费用记录"""
    # 计算费用
    input_cost, output_cost, total_cost = ModelPricing.calculate_cost(
        request_record.model,
        request_record.input_tokens,
        request_record.output_tokens
    )
    
    cost_record = CostRecord(
        user_id=request_record.user_id,
        request_id=request_record.id,
        model=request_record.model,
        input_tokens=request_record.input_tokens,
        output_tokens=request_record.output_tokens,
        input_cost=input_cost,
        output_cost=output_cost,
        cost=total_cost
    )
    
    db.session.add(cost_record)
    
    # 从用户余额中扣除费用（如果启用余额系统）
    user = User.query.get(request_record.user_id)
    if user and user.balance >= total_cost:
        user.balance -= total_cost
    
    return cost_record

def estimate_token_count(text):
    """估算文本的token数量（简单估算）"""
    # 这是一个简单的估算方法，实际应用中可能需要使用tiktoken等库
    # 英文：大约4个字符=1个token
    # 中文：大约1.5个字符=1个token
    
    if not text:
        return 0
    
    # 简单的中英文混合估算
    chinese_chars = len([c for c in text if '\u4e00' <= c <= '\u9fff'])
    other_chars = len(text) - chinese_chars
    
    estimated_tokens = int(chinese_chars / 1.5 + other_chars / 4)
    return max(1, estimated_tokens)  # 至少1个token

def get_model_info(model_name):
    """获取模型信息"""
    pricing = ModelPricing.query.filter_by(model_name=model_name, is_active=True).first()
    if pricing:
        return {
            'model_name': model_name,
            'platform': pricing.platform,
            'input_price_per_1k': pricing.input_price_per_1k,
            'output_price_per_1k': pricing.output_price_per_1k
        }
    
    # 如果没有找到定价信息，返回默认信息
    return {
        'model_name': model_name,
        'platform': 'unknown',
        'input_price_per_1k': 0.01,  # 默认价格
        'output_price_per_1k': 0.03
    }

def validate_user_input(username=None, email=None, password=None):
    """验证用户输入"""
    errors = []
    
    if username is not None:
        if len(username) < 3:
            errors.append("用户名至少需要3个字符")
        if len(username) > 80:
            errors.append("用户名不能超过80个字符")
        if not username.replace('_', '').replace('-', '').isalnum():
            errors.append("用户名只能包含字母、数字、下划线和连字符")
    
    if email is not None:
        import re
        email_pattern = r'^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$'
        if not re.match(email_pattern, email):
            errors.append("邮箱格式不正确")
    
    if password is not None:
        if len(password) < 8:
            errors.append("密码至少需要8个字符")
        if len(password) > 128:
            errors.append("密码不能超过128个字符")
        if not any(c.isalpha() for c in password):
            errors.append("密码必须包含字母")
        if not any(c.isdigit() for c in password):
            errors.append("密码必须包含数字")
    
    return errors

def check_user_exists(username=None, email=None, exclude_user_id=None):
    """检查用户是否已存在"""
    query = User.query
    
    if exclude_user_id:
        query = query.filter(User.id != exclude_user_id)
    
    if username:
        existing_user = query.filter_by(username=username).first()
        if existing_user:
            return True, "用户名已存在"
    
    if email:
        existing_user = query.filter_by(email=email).first()
        if existing_user:
            return True, "邮箱已被注册"
    
    return False, None